#include "ProtoManager.h"
#include "utils/logger.h"

#include <google/protobuf/text_format.h>
#include <iostream>
#include <string>

ProtoManager::ProtoManager(const Operation& preliminaryProto_, const int numberOfWorkers_) 
    : preliminaryProto(preliminaryProto_), numberOfWorkers(numberOfWorkers_)
{
    for (int i = 0; i < numberOfWorkers; ++i)
        stripes.push_back(new internal::TaskList());
}

void ProtoManager::printStripe(const ::google::protobuf::Message &proto, const int workerID) const {
    std::string out;
    google::protobuf::TextFormat::PrintToString(proto, &out);
    std::cout << "Worker ID: " << workerID << std::endl << out << std::endl;
}

std::vector<internal::TaskList*> ProtoManager::getStripes() {
    for (int workerID = 1; workerID <= numberOfWorkers; ++workerID) {
        internal::Operation* o = stripes[workerID-1]->add_task_item()
            ->mutable_final_task()->mutable_source();
        currNodes.push_back(o);

        sliceOperation(preliminaryProto, workerID);
    }

    return stripes;
}

void ProtoManager::sliceOperation(const Operation& op, const int workerID) {
    if (op.has_scan())
        sliceScan(op.scan(), workerID);
    else if (op.has_compute())
        sliceCompute(op.compute(), workerID);
    else if (op.has_filter())
        sliceFilter(op.filter(), workerID);
    else if (op.has_group_by())
        sliceGroupBy(op.group_by(), workerID);
}

void ProtoManager::sliceScan(const ScanOperation& scan, const int workerID) {
    internal::ScanOperation* internalScan = currNodes[workerID-1]->mutable_scan();
    for (int i = 0; i < scan.column_size(); ++i) {
        internalScan->add_column(scan.column(i));
        internalScan->add_type((internal::ScanOperation_Type)scan.type(i));
    }

    //assert(workerID <= scan.number_of_files());
    if (workerID > scan.number_of_files())
        return;

    int firstID = (scan.number_of_files() / numberOfWorkers) * (workerID-1) + 
        std::min(scan.number_of_files() % numberOfWorkers, workerID-1);
    int lastID = firstID + (scan.number_of_files() / numberOfWorkers) + 
        ((scan.number_of_files() % numberOfWorkers) >= workerID);
    for (int i = firstID; i < lastID; ++i)
        internalScan->add_file_number(i);
}

void ProtoManager::copyExpression(const Expression& se, internal::Expression* const de) {
    // operator
    de->set_operator_((internal::Expression_Operator)se.operator_());

    // children
    for (int i = 0; i < se.children_size(); ++i)
        copyExpression(se.children(i), de->add_children());

    if (se.has_constant_int32()) // constant_int32
        de->set_constant_int32(se.constant_int32());

    if (se.has_constant_double()) // constant_double
        de->set_constant_double(se.constant_double());

    if (se.has_constant_bool()) // constant_bool
        de->set_constant_bool(se.constant_bool());

    if (se.has_column_id()) // column_id
        de->set_column_id(se.column_id());
}

void ProtoManager::sliceCompute(const ComputeOperation& compute, const int workerID) {
    internal::ComputeOperation* internalCompute = currNodes[workerID-1]->mutable_compute();
    
    // source
    currNodes[workerID-1] = internalCompute->mutable_source();
    sliceOperation(compute.source(), workerID);

    // expressions
    for (int i = 0; i < compute.expressions_size(); ++i) 
        copyExpression(compute.expressions(i), internalCompute->add_expressions());
}

void ProtoManager::sliceFilter(const FilterOperation& filter, const int workerID) {
    internal::FilterOperation* internalFilter = currNodes[workerID-1]->mutable_filter();

    // source
    currNodes[workerID-1] = internalFilter->mutable_source();
    sliceOperation(filter.source(), workerID);

    // expression
    copyExpression(filter.expression(), internalFilter->mutable_expression());
}

void ProtoManager::sliceGroupBy(const GroupByOperation& group_by, const int workerID) {
    internal::GroupByOperation* internalGroupBy = currNodes[workerID-1]->mutable_group_by();

    // source
    internal::InputOperation* dio = internalGroupBy->mutable_source();
    Operation op;
    GroupByOperation* gbo = op.mutable_group_by();
    gbo->CopyFrom(group_by);
    internal::InputOperation* sio = getResultTypes(op);
    for (int i = 0; i < sio->type_size(); ++i) 
        dio->add_type(sio->type(i));

    // group_by_column_size
    internalGroupBy->set_group_by_column_size(group_by.group_by_column_size());

    // new intermediate task
    internal::BEGroupByOperation* be = stripes[workerID-1]->add_task_item()
        ->mutable_intermediate_task()->mutable_best_effort_group_by();

    // new intermediate task - source
    currNodes[workerID-1] = be->mutable_source();
    sliceOperation(group_by.source(), workerID);

    // new intermediate task - group_by_column
    for (int i = 0; i < group_by.group_by_column_size(); ++i)
        be->add_group_by_column(group_by.group_by_column(i));

    // new intermediate task - aggregations
    for (int i = 0; i < group_by.aggregations_size(); ++i) {
        Aggregation sa = group_by.aggregations(i);
        internal::Aggregation* da = be->add_aggregations();

        // aggregation - type
        da->set_type((internal::Aggregation_Type)sa.type());
        // aggregation - aggregated_column
        if (sa.has_aggregated_column()) 
            da->set_aggregated_column(sa.aggregated_column());
    }
}

internal::InputOperation_Type ProtoManager::getTypeOfExpression(
    const Expression& e, const internal::InputOperation* const io) const 
{
    switch (e.operator_()) {
        case Expression_Operator_CONSTANT:
            if (e.has_constant_int32())
                return internal::InputOperation_Type_INT;
            else if (e.has_constant_double())
                return internal::InputOperation_Type_DOUBLE;
            else if (e.has_constant_bool())
                return internal::InputOperation_Type_BOOL;
            break;
        case Expression_Operator_COLUMN:
            if (e.has_column_id())
                return io->type(e.column_id());
            break;
        case Expression_Operator_IF:
            return getTypeOfExpression(e.children(1), io); // first branch's type
        case Expression_Operator_ADD:
        case Expression_Operator_SUBTRACT:
        case Expression_Operator_MULTIPLY: {
            internal::InputOperation_Type t1 = getTypeOfExpression(e.children(0), io);
            internal::InputOperation_Type t2 = getTypeOfExpression(e.children(1), io);
            if (t1 == ScanOperation_Type_INT && t2 == ScanOperation_Type_INT)
                return internal::InputOperation_Type_INT;
            else // assert: proto is valid
                return internal::InputOperation_Type_DOUBLE;
            }
            break;
        case Expression_Operator_FLOATING_DIVIDE:
        case Expression_Operator_LOG:
            return internal::InputOperation_Type_DOUBLE;
        case Expression_Operator_NEGATE:
            return getTypeOfExpression(e.children(0), io); 
        case Expression_Operator_LOWER:
        case Expression_Operator_GREATER:
        case Expression_Operator_EQUAL:
        case Expression_Operator_NOT_EQUAL:
        case Expression_Operator_NOT:
        case Expression_Operator_OR:
        case Expression_Operator_AND:
            return internal::InputOperation_Type_BOOL;
    };
    CHECK(false, "ProtoManager: invalid proto");
}

internal::InputOperation* ProtoManager::getResultTypes(const Operation& op) const {
    internal::InputOperation* result = new internal::InputOperation();

    if (op.has_filter())
        return getResultTypes(op.filter().source());
    else if (op.has_scan()) {
        ScanOperation scan = op.scan();
        for (int i = 0; i < scan.type_size(); ++i)
            result->add_type((internal::InputOperation_Type)scan.type(i));
    } else if (op.has_compute()) {
        ComputeOperation compute = op.compute();
        internal::InputOperation* io = getResultTypes(compute.source());
        for (int i = 0; i < compute.expressions_size(); ++i)
            result->add_type(getTypeOfExpression(compute.expressions(i), io));
    } else if (op.has_group_by()) {
        GroupByOperation group_by = op.group_by();
        internal::InputOperation* io = getResultTypes(group_by.source());

        // group_by_columns
        for (int i = 0; i < group_by.group_by_column_size(); ++i)
            result->add_type(io->type(group_by.group_by_column(i)));

        // aggregated columns
        for (int i = 0; i < group_by.aggregations_size(); ++i) {
            if (group_by.aggregations(i).type() == Aggregation_Type_COUNT)
                result->add_type(internal::InputOperation_Type_INT);
            else
                result->add_type(io->type(group_by.aggregations(i).aggregated_column()));
        }
    }

    return result;
}
